"""
数据集切分工具
"""
import os
import shutil


class DatasetsCutUtil:
    def __init__(self,
                 read_path: str,
                 train_rate: float = 0.8,
                 train_num: int = 0,
                 cut_type: int = 0,
                 cut_config_type: int = 0):
        """
        切分数据集
        :param read_path:
        :param train_rate:  训练集比例
        :param train_num:  训练集数量
        :param cut_type:  切分类型 0:按配置切分 1:简单切分
        :param cut_config_type:  切分配置类型 0:按比例切分 1:按数量切分
        """
        if not os.path.exists(read_path):
            raise Exception("路径不存在")
        self.read_path = read_path
        self.pic_format = ['bmp', 'jpg', 'png', 'tif', 'gif', 'jpeg']
        # 读取目录下的所有文件
        self.file_list = os.listdir(read_path)
        # 过滤掉非图片文件
        self.file_list = [file for file in self.file_list if file.split('.')[-1] in self.pic_format]
        # 训练集
        self.train_files = []
        # 测试集
        self.test_files = []
        # 切分类型 0:按配置切分 1:简单切分
        self.cut_type = cut_type
        self.cut_config_type = cut_config_type
        self.train_rate = train_rate
        self.train_num = train_num
        # 计算训练集和测试集
        self.calc_train_pic_list()
        # 检查路径是否存在，不存在就创建
        self.check_and_create_dir()
        # 保存训练集和测试集
        self.save_train_and_test_pic()
        # 写出一份yolov5 coco128配置
        self.write_yolov5_coco128()

    def calc_train_pic_list(self):
        """
        计算训练集数量，剩下的就是测试集数量
        :return:
        """
        train_files = []
        test_files = []
        if self.cut_type == 1:
            train_files, test_files = self.simple_cut()
        if self.cut_type == 0:
            train_files, test_files = self.config_cut()
        self.train_files = train_files
        self.test_files = test_files

    def simple_cut(self):
        """
        简单切分
        简单切分就是按照内置的规则切分
        小于300张图片的话 训练集和测试集共用
        如果小于500张，按9:1切分
        如果大于500张，按8:2切分
        :return:  训练集和测试集
        """
        if len(self.file_list) < 300:
            return self.file_list, self.file_list
        elif len(self.file_list) < 500:
            return self.file_list[:int(len(self.file_list) * 0.9)], self.file_list[int(len(self.file_list) * 0.9):]
        return self.file_list[:int(len(self.file_list) * 0.8)], self.file_list[int(len(self.file_list) * 0.8):]

    def config_cut(self):
        """
        按配置切分
        :return:
        """
        if self.cut_config_type == 0:
            return self.config_cut_by_rate()
        if self.cut_config_type == 1:
            return self.config_cut_by_num()

    def config_cut_by_rate(self):
        """
        按比例切分
        :return:
        """
        return self.file_list[:int(len(self.file_list) * self.train_rate)], self.file_list[
                                                                            int(len(self.file_list) * self.train_rate):]

    def config_cut_by_num(self):
        """
        按数量切分
        :return:
        """
        return self.file_list[:self.train_num], self.file_list[self.train_num:]

    def check_and_create_dir(self):
        """
        检查路径是否存在，不存在就创建
        :return:
        """
        yolo_datasets_path = os.path.join(self.read_path, 'yolo_datasets')
        path_list = [
            yolo_datasets_path,
            os.path.join(yolo_datasets_path, 'images'),
            os.path.join(yolo_datasets_path, 'images', 'train'),
            os.path.join(yolo_datasets_path, 'images', 'val'),
            os.path.join(yolo_datasets_path, 'labels'),
            os.path.join(yolo_datasets_path, 'labels', 'train'),
            os.path.join(yolo_datasets_path, 'labels', 'val')
        ]
        for path in path_list:
            if not os.path.exists(path):
                os.makedirs(path)
            else:
                # 清空目录
                shutil.rmtree(path)
                os.makedirs(path)

    def save_train_and_test_pic(self):
        """
        保存训练集和测试集
        :return:
        """
        for file in self.train_files:
            self.copy_file(file, 'train')
        for file in self.test_files:
            self.copy_file(file, 'val')

    def copy_file(self, file, type_):
        """
        复制文件
        :param file: 文件名
        :param type_: 输出类型
        :return:
        """
        source_path = os.path.join(self.read_path, file)
        target_path = os.path.join(self.read_path, 'yolo_datasets', 'images', type_, file)
        # 复制图片
        shutil.copy(source_path, target_path)
        # 复制txt
        txt_file = file.split('.')[0] + '.txt'
        source_txt_path = os.path.join(self.read_path, txt_file)
        target_txt_path = os.path.join(self.read_path, 'yolo_datasets', 'labels', type_, txt_file)
        shutil.copy(source_txt_path, target_txt_path)

    def write_yolov5_coco128(self):
        """
        写出一份yolov5 coco128配置
        :return:  None
        """
        target_path = os.path.join(self.read_path, 'yolo_datasets', 'coco128.yaml')
        with open(target_path, 'w') as f:
            f.write('path: 【数据上传的绝对目录】\n')
            f.write('train: images/train\n')
            f.write('val: images/val\n')
            f.write('nc: 1\n')
            f.write('names: [\'0\']\n')
